# Used to explore the method of using matricies to solve systems of differential equations
# that describe the radioactive decay chains.
# This example takes data from 140228E2 which was analysed with dosDAN, and we plugin the
# initial conditions to get the equations of number of atoms at a given time.
# Rate is simply lambda*N(t).
# This technique is widly used to solve system's of differential equations, but this example
# comes from a paper which was used as a guide to check my solutions.
#      Am. J. Phys. 71 (7) July 2003 - L. Moral and F. Pacheco
# http://users.df.uba.ar/sgil/physics_paper_doc/papers_phys/modern/radiactive_dacay.pdf

import numpy

# Constants and the like
ln2 = 0.69315
l1 = ln2/(3.8235*24) #Rn222
l2 = ln2/(3.1/60.)   #Po218
l3 = ln2/(26.8/60.)  #Pb214
l4 = ln2/(19.9/60.)  #Po214
eff218 = 0.35*0.94 #counting efficiency
eff214 = 0.35*0.94 #counting efficiency

# Rate equations (paper example)
A = numpy.array([[-l1,0,0,0],[l1,-l2,0,0],[0,l2,-l3,0],[0,0,l3,-l4]])

# Initial conditions
n0 = numpy.array([[26449/l1/24.],[0],[0],[0]])

# This is a simple function that prints out matricies in a pretty format
def print_mat(mat):
    if sum(mat.shape)-1 == mat.size:
        # 1D matrix
        for i in range(mat.size):
            print '%.3f'%mat[i]
    else:
        # 2D
        for i in range(len(mat[0,:])):
            s = ''
            for j in range(len(mat[:,i])):
                s+='%.3e\t'%mat[i,j]
            print s
	print "\n"

# Get eigenvalues and vectors
print "Eigenvalues: Eigenvectors"
evalue,evector = numpy.linalg.eig(A)

# Printout results to check what I got for eigenvalue and eigenvector matricies
for i in range(len(evalue)):
   print '%d:'%i, evalue[i], evector[:,i]
print "\n"
print "V="
print_mat(evector)
print "V^{-1}="
print_mat(numpy.linalg.inv(evector))

# How do I order the eigenvalues??? It would seem they should have an order beacuse the
# eigenvalues are basically lambda so they are related to the initial condition vectors.
# Maths is hard, and it probably sholdn't matter what order... (what defines the
# order of roots of a polynomial?)

# This matrix is the solution to the number of atoms in each state as a function of time!!!!1
B = lambda t: numpy.dot(numpy.dot(numpy.dot(evector, numpy.diag(numpy.exp(evalue*t))), numpy.linalg.inv(evector)),n0)
# ... and seems to agree with their results for a large numbers of times tested

# BUT what if we were to put the eigenvectors in the same order as the paper...
evals2 = evalue[::-1]
evects2 = numpy.fliplr(evector)
for i in range(len(evals2)):
   print '%d:'%i, evals2[i], evects2[:,i]
print "\n"
print "V="
print_mat(evects2)
print "V^{-1}="
print_mat(numpy.linalg.inv(evects2))
C = lambda t: numpy.dot(numpy.dot(numpy.dot(evects2, numpy.diag(numpy.exp(evals2*t))), numpy.linalg.inv(evects2)),n0)
# Evalutate B(t) and C(t) over several time and SAME THING! So maths works.
# Again it makes sence it doesn't matter what order as long as the math is correct.

# Lets make some plots to show the results verses some dater.
import matplotlib.pyplot as plt
t = numpy.linspace(0,9,200)
Rn = []
Po218 = []
Po214 = []
Pb210 = []
for i in t:
   tmp = B(i)
   Rn.append(tmp[0])
   Po218.append(tmp[1])
   Po214.append(tmp[3])
   #Pb210.append(1-sum(tmp))

# Ok this solves for the number of atoms as a function of time, actually it gives you a number at "time"
# This probably needs to be numerically integrated and then multiplied by lambda to get the
# number of decays... 

# We want to numerically integrate with resolution of 1/100th of an hour (eventually change units to seconds)
tf = 66 #(hours)
resolution = 3 #hours/point (DAQ)
time = numpy.linspace(0,tf,tf*100+1)
print "Creating N_i(t) data at 1/100th hour resolution"
Po218 = numpy.zeros(time.size-1)
Po214 = numpy.zeros(time.size-1)
Po218_integrated = numpy.zeros(tf/resolution)
Po214_integrated = numpy.zeros(tf/resolution)
for t in time[:-1]:
   tmp = B(t)
   Po218[t] = tmp[1]
   Po214[t] = tmp[3]
   daq_i = int(t/resolution)
   Po218_integrated[daq_i] += Po218[t]/100.
   Po214_integrated[daq_i] += Po214[t]/100.
Po218_integrated = Po218_integrated * l2 * eff218
Po214_integrated = Po214_integrated * l4 * eff214

# We need some data to compare this to, grab the data from file
# Lets quit fitting this damn thing and save it to a pickle...
import pickle
try:
    print "Loaded 140228E2.pkl, a temporary readesc type stage"
    print " If you want to rerun find_peaks and all that delete this file."
    A = pickle.load(open('140228E2.pkl','r'))
except:
    print "No readesc pickle file exists, creating one..."
    import readchn
    print "Opening 140228E2 files..."
    A = readchn.esc_run(run_path = '../data/140228E2.ZIP',verbose=False)
    A.initial_find_peaks()
    A.smart_guess_peaks(0.12)
    A.smart_fit_peaks()
    print " ...Done with smart_fit"
    A.compute_alpha_counts()
    pickle.dump(A,open('140228E2.pkl','w'))

# Make a plot of the Initial value of Rn from dosDAN
leg = []
for alpha in A.alpha_counts.keys():
    if alpha in ['Po210','Po212','Po216']: continue
    print alpha, ": ", A.alpha_counts[alpha].sum()
    plt.plot(A.alpha_counts[alpha],'o')
    leg.append(alpha)
plt.plot(Po214_integrated,'b')
leg.append('Po214 fit')
plt.plot(Po218_integrated,'g')
leg.append('Po218 fit')
plt.ylim(0,1500)
plt.legend(leg)
plt.title('Rn222 initial value = %d atoms'%n0[0])
plt.xlabel('DAQ Interval (3 hours)')
plt.ylabel('Counts (/Interval)')
plt.show()

# Ok, this works now (thanks Jacques), need to write a minimization technique for this
#       (somewhere else)....


